#include <fstream>
#include <iomanip>
#include <iostream>
#include <eigen3/Eigen/Dense>

#include "nav_base.hpp"

class MahonyINS : NavBase{
public:
    bool is_initial = false;
    double sampleFreq = 200;
    double twoKp = 2.0 * 0.2;
    double twoKi = 2.0 * 0.0;
    double integralFBx = 0.0,  integralFBy = 0.0, integralFBz = 0.0;
    Eigen::Quaterniond q_bn;
    Eigen::Vector3d eular;
    Eigen::Matrix3d dcm;

    MahonyINS(){
        q_bn.w() = 1;
        q_bn.x() = q_bn.y() = q_bn.z() = 0;

    }

    void update(Eigen::Vector3d& gyro_rps, Eigen::Vector3d& accel_mpss){
        float recipNorm;
        float halfvx, halfvy, halfvz;
        float halfex, halfey, halfez;
        float qa, qb, qc;

        auto ax = accel_mpss[0];
        auto ay = accel_mpss[1];
        auto az = accel_mpss[2];
        auto gx = gyro_rps[0];
        auto gy = gyro_rps[1];
        auto gz = gyro_rps[2];

        float q0 = q_bn.w();
        float q1 = q_bn.x();
        float q2 = q_bn.y();
        float q3 = q_bn.z();

        // Compute feedback only if accelerometer measurement valid (avoids NaN in accelerometer normalisation)
        if(!((ax == 0.0) && (ay == 0.0) && (az == 0.0))) {

            // Normalise accelerometer measurement
            recipNorm = invSqrt(ax * ax + ay * ay + az * az);
            ax *= recipNorm;
            ay *= recipNorm;
            az *= recipNorm;        

            // Estimated direction of gravity and vector perpendicular to magnetic flux
            halfvx = q1 * q3 - q0 * q2;
            halfvy = q0 * q1 + q2 * q3;
            halfvz = q0 * q0 - 0.5f + q3 * q3;
        
            // Error is sum of cross product between estimated and measured direction of gravity
            halfex = (ay * halfvz - az * halfvy);
            halfey = (az * halfvx - ax * halfvz);
            halfez = (ax * halfvy - ay * halfvx);

            // Compute and apply integral feedback if enabled
            if(twoKi > 0.0f) {
                integralFBx += twoKi * halfex * (1.0f / sampleFreq);	// integral error scaled by Ki
                integralFBy += twoKi * halfey * (1.0f / sampleFreq);
                integralFBz += twoKi * halfez * (1.0f / sampleFreq);
                gx += integralFBx;	// apply integral feedback
                gy += integralFBy;
                gz += integralFBz;
            }
            else {
                integralFBx = 0.0f;	// prevent integral windup
                integralFBy = 0.0f;
                integralFBz = 0.0f;
            }

            // Apply proportional feedback
            gx += twoKp * halfex;
            gy += twoKp * halfey;
            gz += twoKp * halfez;
        }
        
        // Integrate rate of change of quaternion
        gx *= (0.5f * (1.0f / sampleFreq));		// pre-multiply common factors
        gy *= (0.5f * (1.0f / sampleFreq));
        gz *= (0.5f * (1.0f / sampleFreq));
        qa = q0;
        qb = q1;
        qc = q2;
        q0 += (-qb * gx - qc * gy - q3 * gz);
        q1 += (qa * gx + qc * gz - q3 * gy);
        q2 += (qa * gy - qb * gz + q3 * gx);
        q3 += (qa * gz + qb * gy - qc * gx); 
        
        // Normalise quaternion
        recipNorm = invSqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3);
        q0 *= recipNorm;
        q1 *= recipNorm;
        q2 *= recipNorm;
        q3 *= recipNorm;

        q_bn.w() = q0;
        q_bn.x() = q1;
        q_bn.y() = q2;
        q_bn.z() = q3;

        quat2euler(q_bn, eular);
        quat2dcm(q_bn, dcm);
    }


private:
    float invSqrt(float x) {
        float halfx = 0.5f * x;
        float y = x;
        long i = *(long*)&y;
        i = 0x5f3759df - (i>>1);
        y = *(float*)&i;
        y = y * (1.5f - (halfx * y * y));
        return y;
    }
};

int main(int argc, char** argv){

    // 输出文件
    std::ofstream result_fs = std::ofstream("/home/ghowoght/workspace/meg-odom/python/config/17floor/data-17floor-1-10/result/result_ekf.txt");

    MahonyINS mahony_ins;

    Eigen::Vector3d gyro_rps, accel_mpss, odom_mps;
    float timestamp;

    std::ifstream data_fs(
        "/home/ghowoght/workspace/ekf/data/HL_INSPROBE_9_VEL_IMU_ODO.txt");
    data_fs     >> timestamp
                >> gyro_rps[0] >> gyro_rps[1] >> gyro_rps[2]
                >> accel_mpss[0] >> accel_mpss[1] >> accel_mpss[2]
                >> odom_mps[0] >> odom_mps[1];
    odom_mps[0] = (odom_mps[0] + odom_mps[1]) / 2.0;


    Eigen::Vector3d pos;
    pos << 0, 0, 0;
    Eigen::Vector3d gyro_offsets, accel_offsets;
    gyro_offsets << 0, 0, 0;
    accel_offsets << 0, 0, 0;
    int cnt = 0;
    
    while(!data_fs.eof()){
        data_fs     >> timestamp
                >> gyro_rps[0] >> gyro_rps[1] >> gyro_rps[2]
                >> accel_mpss[0] >> accel_mpss[1] >> accel_mpss[2]
                >> odom_mps[0] >> odom_mps[1];
        odom_mps[0] = (odom_mps[0] + odom_mps[1]) / 2.0f;
        odom_mps[1] = odom_mps[2] = 0;

        if(!mahony_ins.is_initial){
            // 对开始时的静止数据求平均得到传感器初始零偏
            gyro_offsets += gyro_rps;
            accel_offsets += accel_mpss;
            cnt++;

            if(cnt == 40000){
                mahony_ins.is_initial = true;
                gyro_offsets /= cnt;
                accel_offsets /= cnt;
                accel_offsets[2] -= -9.79;
            }
        }
        else{
            // 扣除零偏
            gyro_rps -= gyro_offsets;
            accel_mpss -= accel_offsets;
            // 姿态更新
            mahony_ins.update(gyro_rps, accel_mpss);        
            // 航位推算
            auto theta = mahony_ins.eular[0];
            auto dt = 0.005;
            pos[0] += odom_mps[0] * cos(theta) * dt;
            pos[1] += odom_mps[0] * sin(theta) * dt;
            // 保存结果
            result_fs << timestamp << " " 
                << pos[0] << " " << pos[1] << " " << pos[2]
                << std::endl;
        }
    }
    return 0;
}